# coding: utf-8
from netCDF4 import Dataset
import numpy as np
import matplotlib.pyplot as plt
import Gauge


def Deal():
    centers = []
    center2X = GetCenter(
        r'F:/Research/WRF_Output/2008/01/second/20/wrfout_d01_2008-01-18_06_00_00')
    center6X = GetCenter(
        r'F:/Research/WRF_Output/2008/01/second/60/wrfout_d01_2008-01-18_06_00_00')
    center7X = GetCenter(
        r'F:/Research/WRF_Output/2008/01/second/70/wrfout_d01_2008-01-18_06_00_00')
    center18X = GetCenter(
        r'F:/Research/WRF_Output/2008/01/second/180/wrfout_d01_2008-01-18_06_00_00')

    center2X = decumulation(center2X)
    center6X = decumulation(center6X)
    center7X = decumulation(center7X)
    center18X = decumulation(center18X)

    centers.append(center2X)
    centers.append(center6X)
    centers.append(center7X)
    centers.append(center18X)

    gagueCenter, mean = Tool_GetGaugeCenter(
        r"F:/Research/Data/Gauge/Gauge229_07_10_1h.dat")
    centers.append(gagueCenter)
    centers.append(mean)
    labels = ('20', '60', '70', '180', 'gague', 'average')

    time = range(31)
    PlotCenter(time, centers, labels)


def GetCenter(source):
    '''deal with the output of WRF'''

    if source:
        rainfall = Dataset(source)
        rainc = rainfall.variables['RAINC']
        rainnc = rainfall.variables['RAINNC']
        (Time_Number, SN_Number, WE_Number) = rainc.shape
        centerX = SN_Number // 2
        centerY = WE_Number // 2
        center = []
        for t in range(0, Time_Number):
            # get the rainfall data according to time
            rainc_2D = rainc[t, :, :]
            rainnc_2D = rainnc[t, :, :]
            precipitation = rainc_2D[
                centerX, centerY] + rainnc_2D[centerX, centerY]
            center.append(precipitation)
        rainfall.close()
        return center


def PlotCenter(time, centers, labels):
    '''plot rainfall data of center'''
    plt.figure(figsize=(10, 10))
    for i in range(len(centers)):
        plt.plot(time, centers[i], label=labels[i])

    plt.legend(loc='best', fontsize='xx-large')
    plt.show()


def Tool_GetGaugeCenter(dat):
    gauge = Gauge.Read(dat, 229)
    subset2008 = gauge[gauge.index == 2008]
    subset2008 = Gauge.Clean(subset2008)
    gaugeCenter = subset2008.iloc[414:445, 26]
    around = subset2008.iloc[414:445, [22, 93, 175, 126, 94, 178]]
    mean = np.mean(around, axis=1)

    return list(gaugeCenter), mean


def decumulation(seq):
    newSeq = []
    newSeq.append(seq[0])
    for i in range(len(seq) - 1):
        newSeq.append(seq[i + 1] - seq[i])
    return newSeq


if __name__ == '__main__':
    Deal()
